#!/usr/bin/env python

### imports
import sys,os,unittest,getopt
from numpy import array,sum,arange,nan,abs,shape
import numpy as np
sys.path.append("..")
sys.path.append(os.path.join(".","data"))
sys.path.append(os.path.join("..","SpectralMix"))
from SpectralMix import SpectralCluster
from SpectralMix import ClusterBase
from SimulationData import dataScatter, dataScatterLabels
from SimulationData import dataLetters, dataLettersLabels
from SimulationData import dataCircle, dataCircleLabels
from SimulationData import dataNetwork1, dataNetwork1Labels
from GenesExample import geneList, geneLabels, geneDistDict
from SpectralMix.DataManipLib import make_input_graph

### parse inputs 
VERBOSE = False
optlist, args = getopt.getopt(sys.argv[1:], 'v')
for o, a in optlist:
    if o == '-v':
        VERBOSE = True

## class that tests four simulated data sets 
class SpectralClusterTest(unittest.TestCase):

    def setUp(self):
        self.figTypes = ['clustered','original','method']
        self.withPlots = False

    ## scatter plot example
    def testScatterData(self):
        sc = SpectralCluster(dataScatter,k=2,labels=dataScatterLabels,projID='scatter',verbose=VERBOSE)
        self.assertEqual(sc.evaluation['precision'],1.0)
        self.assertEqual(sc.evaluation['recall'],1.0)
        self.assertEqual(sc.evaluation['f1score'],1.0)

    def testFscoreEstimator(self):
        sc = SpectralCluster(dataScatter,k=2,labels=dataScatterLabels,paramEstimator='fscore',projID='scatter',verbose=VERBOSE)
        self.assertEqual(sc.evaluation['precision'],1.0)
        self.assertEqual(sc.evaluation['recall'],1.0)
        self.assertEqual(sc.evaluation['f1score'],1.0)

    ## test circle data example    
    def testCircleData(self):
        sc = SpectralCluster(dataCircle,k=2,paramEstimator='fscore',labels=dataCircleLabels,projID='circle',verbose=VERBOSE,fileType='png')
        #sc.make_plot('eigenspace',labels=sc.origLabels)
        self.assertEqual(sc.evaluation['precision'],1.0)
        self.assertEqual(sc.evaluation['recall'],1.0)
        self.assertEqual(sc.evaluation['f1score'],1.0)
     
    ## test letters data example -- also test the use of alternative ranges for sigma
    def testLetterData(self):
        #sigmaRange = arange(0.001,0.5,0.0001)
        sigmaRange = arange(0.01,5.0,0.01)
        sc = SpectralCluster(dataLetters,k=3,labels=dataLettersLabels,dtype='raw',projID='letters',
                             verbose=VERBOSE,paramEstimator='fscore', sigmaRange=sigmaRange,fileType='png') 
        #print 'letters',sc.sigHat
        #sc.make_plot('eigenspace',labels=sc.origLabels)
        self.assertEqual(sc.evaluation['precision'],1.0)
        self.assertEqual(sc.evaluation['recall'],1.0)
        self.assertEqual(sc.evaluation['f1score'],1.0)

    ## test network data example 
    def testNetworkData(self):
        dn1 = dataNetwork1
        dn1Labels = dataNetwork1Labels
        sc = SpectralCluster(dn1,k=2,dataHeader=dn1.nodes(),labels=dn1Labels,projID='network1',dtype='graph',
                             fileType='png',verbose=VERBOSE)
        self.assertEqual(sc.evaluation['precision'],1.0)
        self.assertEqual(sc.evaluation['recall'],1.0) 
        self.assertEqual(sc.evaluation['f1score'],1.0)
 

    def testSigmaMethods(self):       
        dn1 = dataNetwork1
        dn1Labels = dataNetwork1Labels
        sc = SpectralCluster(dn1,k=2,dataHeader=dn1.nodes(),labels=dn1Labels,projID='network1',dtype='graph',
                             fileType='png',verbose=VERBOSE,paramEstimator='distortion')
       
        sigma = sc.sigHat
        eval1 = sc.evaluation.copy()

        sc = SpectralCluster(dn1,k=2,dataHeader=dn1.nodes(),labels=dn1Labels,projID='network1',dtype='graph',
                             fileType='png',verbose=VERBOSE)
      
        eval2 = sc.evaluation.copy()

        self.assertEqual(eval1['precision'],eval2['precision'])
        self.assertEqual(eval1['recall'],eval2['recall']) 
        self.assertEqual(eval1['f1score'],eval2['f1score'])

    def testLabelingAccuracy(self):
        G = make_input_graph(geneDistDict,geneList)
        k = 2
        
        ## run the spectral clustering
        originalLabels = geneLabels
        sc = SpectralCluster(G,k=k,dataHeader=geneList,labels=geneLabels,projID=os.path.join(".","temp"),
                             dtype='graph',weighted=True,verbose=VERBOSE,sigma=0.14,penalty=True)
        
        ## ensure that the saved header and labes match and original header and labels
        origLabelsBeforeClust = geneLabels
        origLabelsAfterClust = sc.origLabels
        origHeaderBeforeClust = geneList
        origHeaderAfterClust = sc.origDataHeader

        self.assertEqual(len(origLabelsBeforeClust),len(origLabelsAfterClust))
        if len(origLabelsBeforeClust) == len(origLabelsAfterClust):
            for element in range(len(origLabelsAfterClust)):
                self.assertEqual(origLabelsBeforeClust[element],origLabelsAfterClust[element])

        self.assertEqual(len(origHeaderBeforeClust),len(origHeaderAfterClust))
        if len(origHeaderBeforeClust) == len(origHeaderAfterClust):
            for element in range(len(origHeaderAfterClust)):
                self.assertEqual(origHeaderBeforeClust[element],origHeaderAfterClust[element])

        ## make sure the unused genes and indices are consistent
        self.assertEqual(len(sc.unusedGenes), len(sc.unusedIndices))
        self.assertEqual(len(sc.origDataHeader) - len(sc.unusedGenes), len(sc.dataHeader)) 

        ## make sure adjusted data header and labels are correct
        headerCopy = sc.origDataHeader.copy()[sc.usedIndices]
        self.assertEqual(len(headerCopy),len(sc.dataHeader))
        if len(headerCopy) == len(sc.dataHeader):
            for element in range(len(sc.dataHeader)):
                self.assertEqual(headerCopy[element],sc.dataHeader[element])

        labelsCopy = sc.origLabels.copy()[sc.usedIndices]
        self.assertEqual(len(labelsCopy),len(sc.labels))
        if len(labelsCopy) == len(sc.labels):
            for element in range(len(sc.labels)):
                self.assertEqual(labelsCopy[element],sc.labels[element])

    def testInputDtypeConsistancy(self):
        
        dn1 = dataNetwork1
        dn1Labels = dataNetwork1Labels
        sc = SpectralCluster(dn1,k=2,dataHeader=dn1.nodes(),labels=dn1Labels,projID='network1',dtype='graph',
                             fileType='png',verbose=VERBOSE)
       
        affinityMat = sc.aMat.copy()
        eval1 = sc.evaluation.copy()

        sc = SpectralCluster(affinityMat,k=2,dataHeader=dn1.nodes(),labels=dn1Labels,projID='network1',dtype='affinity',
                             fileType='png',verbose=VERBOSE)
      
        eval2 = sc.evaluation.copy()

        self.assertEqual(eval1['precision'],eval2['precision'])
        self.assertEqual(eval1['recall'],eval2['recall']) 
        self.assertEqual(eval1['f1score'],eval2['f1score']) 

    def testMakeInputGraph(self):
        G = make_input_graph(geneDistDict,geneList)
        k = 2
        self.assertEqual(len(geneDistDict),len(G.edges()))

### Run the tests 
if __name__ == '__main__':
    unittest.main()
